from copy import deepcopy
from docplex.mp.model import Model
from QKDNetwork import QKDNetwork

class AddPS:
    def __init__(self, net: QKDNetwork):
        self.net = net
        self.added_node = []

    def run_expr(self, init_ps: list, remove_ilp = False, remove_ilp_rounding = False) -> list:
        self.added_node = init_ps
        ret_rounding = 0
        if(not remove_ilp_rounding):
            ret_rounding = self.rounding()
        ret_ilp = 0
        if(not remove_ilp):
            ret_ilp = self.ilp()
        return [ret_rounding, ret_ilp]
        
    def ilp(self):
        model = Model(name="add_ps_ilp")
        S_dim = self.net.num_nodes + 1
        S = model.binary_var_dict([i for i in range(S_dim)], lb=0, ub=1, name="S")
        e_dim1, e_dim2, e_dim3 = self.net.num_nodes + 1, self.net.num_nodes + 1, self.net.num_sd_pairs
        e = model.binary_var_dict([(i, j, k) for i in range(e_dim1) for j in range(e_dim2) for k in range(e_dim3)], lb=0, ub=1, name="e")
        
        for u in range(self.net.num_nodes):
            for i, sd_item in enumerate(self.net.sd_list):
                flow_expr = model.linear_expr()
                if(sd_item[0] == u):
                    flow_expr -= 1
                if(sd_item[1] == u):
                    flow_expr += 1
                for v in self.net.G.neighbors(u):
                    flow_expr += e[u, v, i]
                    flow_expr -= e[v, u, i]
                model.add_constraint(flow_expr == 0)
        
        for u in range(self.net.num_nodes):
            for v in self.net.G.neighbors(u):
                for i, sd_item in enumerate(self.net.sd_list):
                    model.add_constraint(e[u, v, i] <= S[u] + S[v])
        model.parameters.timelimit.set(60 * 10)
        S_u_sum = model.continuous_var(name="S_u_sum")
        tmp_su_sum = model.linear_expr()
        for i in range(S_dim):
            tmp_su_sum += S[i]
        
        for i in self.added_node:
            model.add_constraint(S[i] == 1)
        
        model.add_constraint(S_u_sum == tmp_su_sum)
        model.minimize(S_u_sum)
        sol = model.solve()
        if sol:
            if model.solve_details.status == 'time limit exceeded':
                return 0
            else:
                return sol.get_value("S_u_sum") - len(self.added_node)
        return 0
        
    def rounding(self):
        model = self.build_rounding_model()
        added_nodes = deepcopy(self.added_node)
        while True:
            sol = model.solve()
            if not sol:
                break
            max_S_u = None
            max_value = 0
            for u in range(self.net.num_nodes):
                S_u_value = sol.get_value("S_{0}".format(u))
                if S_u_value > 0 and S_u_value > max_value and u not in added_nodes:
                    max_S_u = u
                    max_value = S_u_value
            if max_S_u is None:
                break
            added_nodes.append(max_S_u)
            model.add_constraint(model.get_var_by_name(f"S_{max_S_u}") == 1)
            
            # 增加新的PS
            for u in range(self.net.num_nodes):
                S_u_value = sol.get_value("S_{0}".format(u))
                if S_u_value > 0.5 and u not in added_nodes:
                    max_value = S_u_value
                    added_nodes.append(u)
                    model.add_constraint(model.get_var_by_name(f"S_{u}") == 1)
            
        return len(added_nodes) - len(self.added_node)
    
    def build_rounding_model(self) -> Model:
        model = Model(name="add_ps")
        S_dim = self.net.num_nodes + 1
        S = model.continuous_var_dict([i for i in range(S_dim)], lb=0, ub=1, name="S")
        e_dim1, e_dim2, e_dim3 = self.net.num_nodes + 1, self.net.num_nodes + 1, self.net.num_sd_pairs
        e = model.continuous_var_dict([(i, j, k) for i in range(e_dim1) for j in range(e_dim2) for k in range(e_dim3)], lb=0, ub=1, name="e")
        
        for u in range(self.net.num_nodes):
            for i, sd_item in enumerate(self.net.sd_list):
                flow_expr = model.linear_expr()
                if(sd_item[0] == u):
                    flow_expr -= 1
                if(sd_item[1] == u):
                    flow_expr += 1
                for v in self.net.G.neighbors(u):
                    flow_expr += e[u, v, i]
                    flow_expr -= e[v, u, i]
                model.add_constraint(flow_expr == 0)
        
        for u in range(self.net.num_nodes):
            for v in self.net.G.neighbors(u):
                for i, sd_item in enumerate(self.net.sd_list):
                    model.add_constraint(e[u, v, i] <= S[u] + S[v])
        
        for i in self.added_node:
            model.add_constraint(S[i] == 1)
        
        S_u_sum = model.continuous_var(name="S_u_sum")
        tmp_su_sum = model.linear_expr()
        for i in range(S_dim):
            tmp_su_sum += S[i]
        
        model.add_constraint(S_u_sum == tmp_su_sum)
        model.minimize(S_u_sum)
        
        return model
        